#ifndef SHADERALLOCATOR_H
#define SHADERALLOCATOR_H

#include "BucketAllocator.h"
#include "CryMemoryAllocator.h"
#include "IMemory.h"

#if defined(USE_GLOBAL_BUCKET_ALLOCATOR) && (defined(XENON) || defined(PS3))
#define SHADER_ALLOCATOR_USE_BUCKETS
#endif

#ifdef SHADER_ALLOCATOR_USE_BUCKETS

#define SHADER_BUCKET_ALLOCATOR_SIZE (12 * 1024 * 1024)
typedef BucketAllocator<BucketAllocatorDetail::DefaultTraits<SHADER_BUCKET_ALLOCATOR_SIZE, BucketAllocatorDetail::SyncPolicyLocked, false> > ShaderBucketAllocator;

#else

typedef cry_crt_node_allocator ShaderBucketAllocator;

#endif

extern ShaderBucketAllocator g_shaderBucketAllocator;
extern IGeneralMemoryHeap* g_shaderGeneralHeap;

template <class T>
class STLShaderAllocator
{
public:
	typedef size_t    size_type;
	typedef ptrdiff_t difference_type;
	typedef T*        pointer;
	typedef const T*  const_pointer;
	typedef T&        reference;
	typedef const T&  const_reference;
	typedef T         value_type;

	template <class U> struct rebind
	{
		typedef STLShaderAllocator<U> other;
	};

	STLShaderAllocator() throw() { }
	STLShaderAllocator(const STLShaderAllocator&) throw() { }
	template <class U> STLShaderAllocator(const STLShaderAllocator<U>&) throw() { }

	pointer address(reference x) const
	{
		return &x;
	}

	const_pointer address(const_reference x) const
	{
		return &x;
	}

	pointer allocate(size_type n = 1, const void* hint = 0)
	{
		MEMREPLAY_SCOPE(EMemReplayAllocClass::C_UserPointer, EMemReplayUserPointerClass::C_STL);
		
		pointer ret = NULL;
		
		(void)hint;
		size_t sz = std::max<size_type>(n * sizeof(T), 1);
		if (sz <= ShaderBucketAllocator::MaxSize)
		{
			ret = static_cast<pointer>(g_shaderBucketAllocator.allocate(sz));
		}
		else
		{
			ret = static_cast<pointer>(g_shaderGeneralHeap->Malloc(sz, NULL));
		}
		
		MEMREPLAY_SCOPE_ALLOC(ret, n * sizeof(T), 0);
		
		return ret;
	}

	void deallocate(pointer p, size_type n = 1)
	{
		MEMREPLAY_SCOPE(EMemReplayAllocClass::C_UserPointer, EMemReplayUserPointerClass::C_STL);
		
		(void)n;
		if (p)
		{
			if (!g_shaderGeneralHeap->Free(p))
				g_shaderBucketAllocator.deallocate(p);
		}
		
		MEMREPLAY_SCOPE_FREE(p);
	}

	size_type max_size() const throw()
	{
		return INT_MAX;
	}

	void construct(pointer p, const T& val)
	{
		new(static_cast<void*>(p)) T(val);
	}

	void construct(pointer p)
	{
		new(static_cast<void*>(p)) T();
	}

	void destroy(pointer p)
	{
		p->~T();
	}

	pointer new_pointer()
	{
		return new(allocate()) T();
	}

	pointer new_pointer(const T& val)
	{
		return new(allocate()) T(val);
	}

	void delete_pointer(pointer p)
	{
		p->~T();
		deallocate(p);
	}

	bool operator==(const STLShaderAllocator&) {return true;}
	bool operator!=(const STLShaderAllocator&) {return false;}

	static void GetMemoryUsage( ICrySizer *pSizer ) { }
};

#endif
